Add additive symbolic regression (stagewise & backfitting) - #11

Adversarial probing of jaxsr.additive surfaced two silent footguns at the
new public API boundary:

- Non-finite inputs (NaN/inf in X or y) silently propagated to NaN
  predictions. fit() now raises ValueError.
- predict() silently accepted inputs with the wrong number of features
  (using/ignoring columns) and returned a plausible-but-wrong result.
  AdditiveSymbolicModel.predict now validates the feature count.

Adds regression tests for both, plus tests for single-feature data,
tiny-sample fits, and reproducibility under a fixed random_state.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

A second, correctness-focused round of adversarial testing confirmed the
model's semantics are sound (predict matches the internal training
prediction to ~1e-9; to_expression matches predict to ~1e-7; refit loss is
exactly monotone; early-stopping rollback state is consistent). The one
real gap was persistence: like every jaxsr estimator, additive models are
not picklable (basis-function closures), and no idiomatic save/load existed.

Add JSON save/load via _state_dict/_from_dict, delegating each term to the
existing SymbolicRegressor serialization. Covers refit, non-refit, and
early-stopping models. Adds a round-trip test and documents the pattern in
the guide.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

A behavioral train/test study surfaced a real defect: with
include_transcendental=True on symmetric-domain data, a stage could select a
basis like log(x0) that is invalid for negative x0. The base package zeroes
non-finite design-matrix entries during fitting (so the corrupted column can
be selected) but predict() recomputes it raw, yielding NaN across all
predictions. This is an inherited base-package fit/predict inconsistency,
exposed by additive's opt-in transcendental support.

_fit_stage now checks each fitted term for finite predictions on the
training data and, if a transcendental/ratio basis produced NaN, refits the
stage without those bases and warns. The ensemble can no longer contain a
NaN-producing term. Adds a regression test.

The broader behavioral study otherwise validated the design: stagewise with
max_complexity=1 recovers a 4-term target (R2=0.99) a matching single-term
fit cannot (R2=0.30); interactions recover exactly; early stopping helps on
noisy data; collinear features are handled.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Metamorphic testing found that scaling the target by 1e-6 collapsed the fit
(R2 0.055) while 1e6 was fine. Root cause: refit_ols fit the intercept by
augmenting Phi with a ones-column (magnitude 1) alongside term columns whose
scale tracks y (~1e-6 for tiny targets). That ~1e6 condition number is
unresolvable in JAX's default float32, so lstsq returned garbage.

Fit the intercept by centering Phi and y instead of augmenting -- the
standard, numerically stable approach that keeps the design matrix on a
single scale. Predictions are unchanged for well-scaled data; tiny/huge
targets now recover correctly.

All 15 metamorphic properties now hold (scale/shift/negation equivariance
to ~1e-7, row-permutation and sample-duplication invariance, statistical
consistency, monotone R2 in n_terms, feature-scaling stability). Adds a
scale-equivariance regression test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

…mples

- README: Features bullet + "Additive Symbolic Regression" section with a
  runnable example; note the standalone example script.
- docs: feature bullet in intro; new API page (docs/api/additive.rst) wired
  into the API index; example notebook
  (docs/examples/additive_symbolic_regression.ipynb, executed) wired into
  the TOC alongside the existing guide.
- skill: Quick Reference block, decision-tree entry, and guides/additive.md
  in .claude/skills/jaxsr; re-synced to src/jaxsr/skill.
- CLAUDE.md: mark additive guide as covered.

All snippets verified against the actual API; README example executes
(R2=0.9985); notebook executes without errors.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Implements robust and quantile symbolic regression -- the capability that
distinguishes additive SR from ordinary least-squares symbolic regression.
Each weak learner fits the negative gradient of the chosen loss, so the
ensemble can target objectives OLS selection cannot.

losses.py:
- AbsoluteError (median init, sign gradient), HuberLoss(delta) (robust,
  quadratic near zero / linear tails), QuantileLoss(quantile) (pinball loss,
  empirical-quantile init). Loss-optimal constant initialisation per loss.
- Parameterized losses are passable by name (defaults) or as instances
  (QuantileLoss(0.9), HuberLoss(2.0)); to_config/loss_from_config make them
  round-trip through save/load.

stagewise.py:
- Non-squared losses are fit by gradient boosting with a per-stage line
  search (scipy minimize_scalar) that picks the loss-optimal step size,
  shrunk by learning_rate.
- refit_coefficients (OLS) targets squared error, so it is auto-disabled
  with a warning for non-squared losses. The squared-error path is unchanged.
- save/load serialises the loss faithfully (name or {name, params}).

Verified: under 8% heavy contamination, MAE-vs-clean is 2.76 (squared) ->
0.14 (huber) -> 0.05 (absolute); quantile coverage tracks the target
(0.10/0.51/0.89 for q=0.1/0.5/0.9). Adds tests for the registry, gradients,
robustness, coverage, the non-squared refit warning, and parameterized-loss
save/load. Updates README, docs guide, skill, and the example notebook.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Phase 1 of item #5: BackfittingSymbolicRegressor now maintains a fixed set of
terms and revises each one across sweeps instead of freezing them. Each sweep
removes a term, re-discovers its expression on the partial residual (via
fit_symbolic), reinserts it, and re-solves all coefficients by OLS. It is
warm-started from a stagewise fit, keeps the best-loss iterate (structure
re-discovery is a heuristic), converges on `tol`, and supports squared error
only (non-squared raises NotImplementedError).

Refactor: shared machinery (fitted-attribute accessors, predict/score/
to_expression, JSON save/load, repr) is factored into base._BaseAdditiveRegressor;
StagewiseSymbolicRegressor and BackfittingSymbolicRegressor both inherit it.
The stagewise fitting logic and squared-error path are unchanged.

Honest benchmark (the reason to gate this feature): for squared error,
backfitting *matches* stagewise+refit rather than beating it (mean test-R2
deltas ~0 across simple/multi-effect/correlated/interaction targets; sweeps
improve warm-start train MSE by ~0%). This is expected -- terms are linear in
their bases and coefficients are refit jointly, so the fit over the union of
discovered bases is largely partition-invariant. Shipped and documented as the
GAM-style "revisable terms" variant and the foundation for a future Bayesian
(BART/iBART) backfitting variant, with users steered to the stagewise regressor
by default.

Adds backfitting tests (fit/recover, parity with stagewise, NotImplementedError
for non-squared, param/input validation, save/load). Updates README, guide,
skill, API docs, and the example notebook.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

…tuck

Earlier benchmark examples were too easy and wrongly suggested backfitting
merely matches stagewise+refit. Instrumented adversarial search shows
otherwise: backfitting starts from the stagewise+refit fit and keeps the best
iterate, so it is never worse on training, and it genuinely improves the fit
in the regime it is designed for -- small per-term budgets (max_complexity=1,
GAM-style single-basis terms) with collinear features, where greedy forward
selection locks into a suboptimal basis set that a single forward pass cannot
undo. Re-discovery changes the selected-basis union and escapes it (measured up
to ~+0.04 train / +0.06 test R^2; training-loss cut ~49% in the best case).

The prior "typically matches" claim held only when greedy already found a
sufficient basis set (generous per-term budgets), which the easy examples did.

- Replace the misleading parity test with the guaranteed invariant
  (backfitting train score >= stagewise+refit) plus a collinear single-basis
  regression test.
- Correct the framing in the guide, README, and skill to state precisely when
  backfitting helps vs matches.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Quantifies how stable the discovered *structure* is, which the collinear-regime
finding showed is the real open question. bootstrap_additive refits the model on
bootstrap resamples and returns per-basis inclusion probabilities (a frequentist
proxy for posterior inclusion probability) plus the fitted ensemble;
bootstrap_predict_additive turns that ensemble into predictive intervals that
reflect structural variability, not just coefficient noise. Both work for the
stagewise and backfitting regressors and reuse each estimator's own fitting
machinery (clone via get_params).

Empirically it reproduces the Phase-1 signal: identifiable data gives crisp
inclusion probabilities (~1.0 for the true bases), collinear data gives diffuse
ones (e.g. x2=1.00, x1=0.60, x0=0.52) -- the honest indicator that no single
expression is determined. This also serves as a cheap decision gate for whether
a full Bayesian treatment is worth it.

Follows jaxsr's bootstrap_* convention (standalone functions returning plain
dicts). Adds tests (structure/reproducibility, instability detection, predictive
intervals, backfitting support, input validation), a guide section, README and
skill notes, an API page, and an executed notebook section.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Document the one-line positioning up front: JAXSR is a linear method over a
fixed basis library (selects a sparse combination of user-supplied basis
functions), not a free-composition equation discoverer. Adds a "Scope" section
to the guide and a short callout in the README explaining that the limit is
discovery, not representation -- the linear-in-basis model fits targets like
exp(x0*x1) perfectly once the exact term is in the library (add_custom), it just
cannot discover which composition is needed. Points users to PySR/Operon/
AI-Feynman for compositional discovery. Skill guide re-synced.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

SymbolicRegressor could return NaN from predict() when a basis function that is
non-finite on the training data (e.g. log(x) or sqrt(x) with x <= 0) was
selected. The design matrix zeroed only the non-finite *elements*, leaving a
corrupted-but-nonzero column that greedy search could still pick to fill a term
slot; predict() then re-evaluated that basis from scratch and produced NaN. This
was surfaced live by the Operon benchmark (a log(x0) term turned an otherwise
correct fit into R2 = -inf).

Fix: zero the whole invalid column so the selection math stays finite, then drop
any non-finite-on-training term that still gets selected and refit the remaining
terms by OLS (recomputing AIC/BIC/AICc), before constraints are applied. The
warning now says "Excluding ... they will not be selected", matching behavior.

Adds a regression test. Note: this addresses terms non-finite on *training*
data; a basis that is finite on training but diverges out-of-domain at predict
(e.g. exp(x0/x1) near x1=0) is a separate issue and not covered here. The same
latent pattern exists in classifier.py and is left for a follow-up.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

A) SymbolicClassifier had the same latent bug as the regressor: a basis
non-finite on the training data (e.g. log(x) with x <= 0) could be selected
and make predict_proba() return NaN. Fixed by zeroing the whole invalid column
during fitting and dropping any that still get selected. Because the column was
zero during fitting it contributed nothing to the logits, so the drop is exact
(in-domain predictions unchanged); information criteria are recomputed from the
unchanged negative log-likelihood. Handles binary and one-vs-rest multiclass.

B) SymbolicRegressor now prunes numerically negligible terms after fitting
(new prune_tol=1e-6 parameter). A term whose contribution |coef|*||basis|| is
below prune_tol times the largest term's is dropped and the rest refit. This
removes spuriously selected bases that are finite on training but diverge out
of domain at predict (e.g. exp(x0/x1) near x1=0, which previously turned a
correct exp(x0*x1) fit into R2=-inf on a single test point). The real vs noise
gap is ~8 orders of magnitude, so real terms are never touched; prune_tol=0
disables it. Post-selection cleanups are skipped for parametric libraries,
whose terms are only meaningful after their internal parameters are optimised.

Refactors the regressor's OLS refit into a shared _refit_subset helper. Adds
regression tests for both classes; updates the sklearn get_params test for the
new parameter. Full suite: 553 passed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

Prototype that partially lifts the fixed-library ceiling. Instead of
enumerating a combinatorial depth-d composition space up front, it grows the
basis library lazily along the residual: fit a sparse model, take the residual,
compose the currently useful terms (selected terms + features) with a small
operator set (unary funcs, products, ratios) one layer, screen candidates
(drop non-finite, dedup, keep top-beam_width by residual correlation), add
survivors, refit, repeat. Effective composition depth = number of rounds. This
is Fast Function Extraction / symbolic feature construction -- a deterministic,
bounded cousin of genetic programming.

The result is an ordinary fitted SymbolicRegressor over the grown library, so
it reuses predict/expression/scoring and the base regressor's non-finite-basis
guard and negligible-term pruning.

Measured: on compositional targets a flat library misses, it substantially
wins (exp(x0*x1): R2 1.00 vs 0.70; x0*sin(x1): 0.9995 vs 0.958) and is
competitive with Operon (matched it on x0*sin(x1); 0.998 vs 0.999 on a
rational). It recovers exp((x0)*(x1)) exactly. Honest limits: cost grows with
beam_width/n_expansions/n_features, it won't match a mature GP on hard targets,
and the composed closures make the model non-serialisable.

Adds tests, a guide section (cross-referenced from the Scope section), README
and skill notes, and an API page.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz

- examples/additive_symbolic_regression.py: expanded from a minimal stagewise
  demo into six runnable sections (stagewise, robust Huber/absolute, quantile
  coverage, bootstrap structural uncertainty, backfitting, recursive expansion).
- docs/examples/additive_symbolic_regression.ipynb: added an executed
  "Recursive basis expansion" section (recovers exp((x0)*(x1)) at R2=1.0 vs a
  flat library's 0.72).
- New skill template templates/additive-regression.py covering all six
  workflows, registered in the SKILL.md templates table; skill re-synced.

Docs/examples only; all snippets verified to run.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QyBMabiyYFLSj8bkJ9MPdz